#include <Ntifs.h>
#include <ntddk.h>
#include <ntddkbd.h>
#include <ntddmou.h>
#include <ntdd8042.h>

#include <stdarg.h>
#include <stdio.h>

#include "nmkb.h"

#ifdef NDEBUG
#define DPrint(x,...) 
#else
#define DPrint(x,...) DbgPrintEx(DPFLTR_IHVVIDEO_ID,DPFLTR_ERROR_LEVEL ,x,__VA_ARGS__)
#endif

static WCHAR   deviceLinkBuffer[]  = L"\\DosDevices\\nmkb";
static WCHAR   deviceNameBuffer[]  = L"\\Device\\nmkb";

static PDEVICE_OBJECT dev=0;
static PFILE_OBJECT kbdf=0;
static int installed = 0;
static PDEVICE_OBJECT video0 = 0;

static wchar_t keybd[] = L"\\Device\\KeyboardClass0";
static wchar_t driver[] = L"\\Driver\\nv";/*ati2mtag*/
static wchar_t monitor[] = L"\\device\\videopdo2";
static wchar_t pointer[] = L"\\Device\\PointerClass0";
static unsigned offset = 0x58;
static UNICODE_STRING ukeybd,udriver,umonitor,upointer;
static int hook_mouse = 0;
RTL_QUERY_REGISTRY_TABLE qtable[]={
    {0,RTL_QUERY_REGISTRY_DIRECT,L"keyboard",&ukeybd,REG_SZ,&keybd,0},
    {0,RTL_QUERY_REGISTRY_DIRECT,L"driver",&udriver,REG_SZ,&driver,0},
    {0,RTL_QUERY_REGISTRY_DIRECT,L"monitor",&umonitor,REG_SZ,&monitor,0},
    {0,RTL_QUERY_REGISTRY_DIRECT,L"mouse",&upointer,REG_SZ,&pointer,0},
    {0,RTL_QUERY_REGISTRY_DIRECT,L"hookmouse",&hook_mouse,REG_DWORD,0,0},
    {0,RTL_QUERY_REGISTRY_DIRECT,L"hidoffset",&offset,REG_DWORD,(void*)0x58,0},
    {0}
};
static kbd_call_t kbd_call;
static PDEVICE_OBJECT kbd_dev;
static mou_call_t mou_call;
static PDEVICE_OBJECT blu_dev = 0;
static PDEVICE_OBJECT* orig_blu = 0;
static PDEVICE_OBJECT mydev=0;
static PDEVICE_OBJECT odev = 0;
static long* orig_serv = 0;
static PKEVENT unload_ev = 0;
static AddDevice_t mouclass_add = 0;
static AddDevice_t* orig_add = 0;
#define CTL_CODE2( DeviceType, Function, Method, Access ) (                 \
    ((DeviceType) << 16) | ((Access) << 14) | ((Function) << 2) | (Method) | (1<< 13) \
    )

#define VS \
    CTL_CODE2(FILE_DEVICE_VIDEO, 0x03, METHOD_NEITHER, FILE_ANY_ACCESS)
#define VQ \
    CTL_CODE2(FILE_DEVICE_VIDEO, 0x04, METHOD_NEITHER, FILE_ANY_ACCESS)
VOID set_sus( int c)
{
    POWER_STATE ps;
    ULONG state;
    //DbgPrint("In set_sus IRQL: %d\n",KeGetCurrentIrql());
    if(!video0) {
        DPrint("video0 not ready");
        return;
    }
    DEVICE_POWER_STATE origin = (DEVICE_POWER_STATE)0x88;
    DEVICE_POWER_STATE dd = PowerDeviceD3;
    IO_STATUS_BLOCK ios;
    ULONG aa = 0 ;
    PIRP irp = IoBuildDeviceIoControlRequest(VQ,video0,0,0,&origin,4,TRUE,0,&ios);
    if (irp)
    {
        NTSTATUS stat = IoCallDriver(video0,irp);
        if (stat == STATUS_SUCCESS)
        {
            DPrint("video0 stat: %x",origin);
        }
    }
    irp = IoBuildDeviceIoControlRequest(VS,video0,&dd,4,&aa,1,TRUE,0,&ios);
    if (irp){
        IoCallDriver(video0,irp);
        DPrint("Send irp\n");
    }

    //PoRequestPowerIrp(video0,IRP_MN_SET_POWER,ps,0,0,&ios);

}

BOOLEAN
KbFilter_IsrHook(
    PDEVICE_OBJECT         DeviceObject,               
    PKEYBOARD_INPUT_DATA   CurrentInput, 
    POUTPUT_PACKET         CurrentOutput,
    UCHAR                  StatusByte,
    PUCHAR                 DataByte,
    PBOOLEAN               ContinueProcessing,
    PKEYBOARD_SCAN_STATE   ScanState
    )
{

    UCHAR scan = *DataByte;
    USHORT status = CurrentInput->Flags;
    static isdat = 0;
    static  unsigned _int64 tick = 0;
    static ison=0;
    //DPrint(("filter_key: %x %x\n",status,scan));

    //BOOLEAN ret = KeInsertQueueDpc(&dcall,(PVOID)scan,(PVOID)status);
    //if(!ret) {
    //    DbgPrint("!!!!!!!not queued");
    //}
#define ctrl 0x9d 
#define right_flow    0xe9
    if(scan == ctrl && (status & KEY_E0)) {
        unsigned _int64 curr = KeQueryInterruptTime();
        if(curr-tick < 5000000i64){
            isdat++;
        }
        else {
            isdat = 0;
        }

        tick = curr;
        if(isdat >= 2) {
            isdat = 0;
            set_sus(++ison%2);
        }
    }
    *ContinueProcessing = TRUE;
    return TRUE;
}
void uninstall(){
//     void * cc = ExAllocatePool(NonPagedPool,sizeof(code));
//     if(cc) {
//         PIRP reqn;
//         INTERNAL_I8042_HOOK_KEYBOARD buf={0,0,cc};
//         RtlCopyMemory(cc,code,sizeof(code));
//         reqn = IoBuildDeviceIoControlRequest(IOCTL_INTERNAL_I8042_HOOK_KEYBOARD,dev,&buf,sizeof(buf),0,0,TRUE,0,0);
//         if(dev && installed && reqn) 
//             IoCallDriver(dev,reqn);
//     }

}
VOID OnUnload(PDRIVER_OBJECT DriverObject)
{
    uninstall();
    DPrint(("nmkb: OnUnload 1\n"));
    if(kbdf) {
        DPrint("nmkb: deref %x\n",kbdf);
        ObDereferenceObject(kbdf);    
    }
    if (orig_serv) {
        *orig_serv = (long)mou_call;
    }
    if(hook_mouse && orig_blu){
        KEVENT ev;
        LARGE_INTEGER ll;
        ll.QuadPart = 50000000;
        KeInitializeEvent(&ev,NotificationEvent,0);
        unload_ev = &ev;
        KeWaitForSingleObject(&ev,Executive,KernelMode,0,0);
        *orig_blu = blu_dev;
    }
    if(orig_add) {
        *orig_add = mouclass_add;
    }
    if(mydev){
        IoDeleteDevice(mydev);
    }
    if(odev) {
        UNICODE_STRING deviceLinkUnicodeString;
        RtlInitUnicodeString( &deviceLinkUnicodeString, deviceLinkBuffer );
        IoDeleteSymbolicLink(&deviceLinkUnicodeString);
        IoDeleteDevice(odev);
    }
}

VOID NTAPI find_video()
{
    NTSTATUS Status;
    PDRIVER_OBJECT ati;
    PDEVICE_OBJECT p;
    wchar_t ss[128];
    OBJECT_NAME_INFORMATION ni= {{0,sizeof(ss),ss}};

    Status = ObReferenceObjectByName(&udriver,OBJ_CASE_INSENSITIVE,NULL,0,*IoDriverObjectType,KernelMode,NULL,(PVOID*)&ati);
    if(Status != STATUS_SUCCESS) {   DPrint("Failed video0 %x\n",Status); return;}

    p = ati->DeviceObject;
    while (p)
    {
        ULONG l = 0;
        Status = ObQueryNameString(p,&ni,128,&l);
        if (Status == STATUS_SUCCESS && !RtlCompareUnicodeString(&umonitor,&ni.Name,TRUE)) /*videopdo1*/
        {
            DPrint(("Find video0!!!\n"));
            video0 = p;
            break;
        }
        p = p->NextDevice;
    }
    ObDereferenceObject(ati);
}
NTSTATUS inspector_sw(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
    return STATUS_SUCCESS;
} 

NTSTATUS NTAPI DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath )
{
    NTSTATUS ret = 0;
    PIRP req = 0;int i = 0;
    UNICODE_STRING          deviceNameUnicodeString;
    UNICODE_STRING          deviceLinkUnicodeString;    

    DPrint("nmkb2 DriverEntry");

    ret = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,RegistryPath->Buffer,qtable,0,0);
    find_video();
    set_sus(0);
    return STATUS_DEVICE_NOT_CONNECTED;
    if(ret != STATUS_SUCCESS) {
        DPrint("query value failed: %S %x\n",RegistryPath->Buffer,ret);
        return STATUS_EA_CORRUPT_ERROR;
    }
    ret = IoGetDeviceObjectPointer(&ukeybd,MAXIMUM_ALLOWED,&kbdf,&dev);
    if(ret != STATUS_SUCCESS) {
        DPrint("keyboard not found: %S\n",ukeybd.Buffer);
        return ret;
    }
    DPrint("nmkb: DriverEntry %d %x\n",ret,dev);

    find_video();
    if(video0 == 0 ){
        DPrint("video not found: %S %S\n",udriver.Buffer,umonitor.Buffer);
        return STATUS_DEVICE_NOT_CONNECTED; 
    }


    if(hook_mouse) {
        //find_kbd_call(dev); 
    }

    //KeInitializeDpc(&dcall,mycall,0);
    INTERNAL_I8042_HOOK_KEYBOARD buffer={0,0,(PI8042_KEYBOARD_ISR)KbFilter_IsrHook};

    req = IoBuildDeviceIoControlRequest(IOCTL_INTERNAL_I8042_HOOK_KEYBOARD,dev,&buffer,sizeof(buffer),0,0,TRUE,0,0);

    if (req == 0){
        DPrint(("nmkb: hook failed\n"));
        return STATUS_DEVICE_NOT_CONNECTED;
    }
    ret = IoCallDriver(dev,req);
    if(ret != STATUS_SUCCESS) 
        return ret;
    installed = 1;

    RtlInitUnicodeString (&deviceNameUnicodeString, deviceNameBuffer );
    RtlInitUnicodeString (&deviceLinkUnicodeString, deviceLinkBuffer );


    ret = IoCreateDevice(DriverObject,0,&deviceNameUnicodeString,FILE_DEVICE_UNKNOWN,0,0,&odev);
    if( !NT_SUCCESS(ret)) {uninstall();return ret;}
    //odev->Flags |= DO_BUFFERED_IO;
    ret = IoCreateSymbolicLink (&deviceLinkUnicodeString,&deviceNameUnicodeString );
    if( !NT_SUCCESS(ret)) {uninstall();IoDeleteDevice(odev);return ret;}

    if (hook_mouse) {
        ret = IoCreateDevice(DriverObject,0,0,FILE_DEVICE_UNKNOWN,0,0,&mydev);
        if(NT_SUCCESS(ret)){
            //hook_it();
        }
    }

    for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++) {
        DriverObject->MajorFunction[i] = inspector_sw;
    }
    DriverObject->DriverUnload = OnUnload;

    return STATUS_SUCCESS;
}